
#pragma once
#include <pcl/registration/ndt.h>
#include "registration_interface.hpp"

namespace TESTICP
{
    class NDTRegistration : public RegistrationInterface
    {
    public:
        NDTRegistration(const YAML::Node& node);
        NDTRegistration(float res, float step_size, float trans_eps, int max_iter);

        void setInputTarget(const CloudPtr& target) override;
        void setInputSource(const CloudPtr& source) override;
        bool align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose = Eigen::Matrix4f::Identity()) override;

        inline Eigen::Matrix4f getFinalTransformation() { return ndt_ptr_->getFinalTransformation(); }
        inline float getFitnessScore() { return ndt_ptr_->getFitnessScore(); }

    private:
        bool SetRegistrationParam(float res, float step_size, float trans_eps, int max_iter);

    private:
        pcl::NormalDistributionsTransform<PointType, PointType>::Ptr ndt_ptr_;
    };

    NDTRegistration::NDTRegistration(const YAML::Node& node)
        : ndt_ptr_(new pcl::NormalDistributionsTransform<PointType, PointType>()) {

        float res = node["res"].as<float>();
        float step_size = node["step_size"].as<float>();
        float trans_eps = node["trans_eps"].as<float>();
        int max_iter = node["max_iter"].as<int>();

        SetRegistrationParam(res, step_size, trans_eps, max_iter);
    }

    NDTRegistration::NDTRegistration(float res, float step_size, float trans_eps, int max_iter)
        : ndt_ptr_(new pcl::NormalDistributionsTransform<PointType, PointType>()) {

        SetRegistrationParam(res, step_size, trans_eps, max_iter);
    }

    bool NDTRegistration::SetRegistrationParam(float res, float step_size, float trans_eps, int max_iter) {
        ndt_ptr_->setResolution(res);
        ndt_ptr_->setStepSize(step_size);
        ndt_ptr_->setTransformationEpsilon(trans_eps);
        ndt_ptr_->setMaximumIterations(max_iter);

        std::cout << "NDT params:" << std::endl
            << "res: " << res << ", "
            << "step_size: " << step_size << ", "
            << "trans_eps: " << trans_eps << ", "
            << "max_iter: " << max_iter
            << std::endl
            << std::endl;

        return true;
    }

    void NDTRegistration::setInputTarget(const CloudPtr& target) {
        ndt_ptr_->setInputTarget(target);
    }

    void NDTRegistration::setInputSource(const CloudPtr& source) {
        ndt_ptr_->setInputSource(source);
    }

    bool NDTRegistration::align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose) {
        ndt_ptr_->align(output_ptr, predict_pose);
        return true;
    }

}
